package com.llmops.agentsflex.demo;

import com.agentsflex.core.chain.*;
import com.agentsflex.core.chain.node.EndNode;
import com.agentsflex.core.chain.node.LlmNode;
import com.agentsflex.core.chain.node.StartNode;
import com.agentsflex.core.prompt.template.TextPromptTemplate;
import com.agentsflex.core.util.Maps;
import com.agentsflex.llm.deepseek.DeepseekConfig;
import com.agentsflex.llm.deepseek.DeepseekLlm;

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;

public class Test {
    public static void main(String[] args) {
        StartNode startNode = new StartNode();
        startNode.setId("start");

        LlmNode llmNode = new LlmNode();
        llmNode.setId("llm");
        llmNode.setSystemPrompt("回答用户的问题");
        llmNode.setUserPrompt("{{startInput}}");
        DeepseekConfig config = new DeepseekConfig();
        config.setEndpoint("https://api.deepseek.com/v1");
        config.setApiKey("sk-0c3e6a9f9f7a4e63bb855290c544183c");
        config.setModel("deepseek-chat");
        config.setDebug(true);
        llmNode.setLlm(new DeepseekLlm(config));
        Parameter parameter1 = new Parameter();
        parameter1.setName("startInput");
        parameter1.setRef("start.input");
        parameter1.setDataType(DataType.String);
        parameter1.setRefType(RefType.REF);
        parameter1.setRequired(true);
        llmNode.addInputParameter(parameter1);

        Parameter parameter = new Parameter();
        parameter.setName("input");
        parameter.setDataType(DataType.String);
        parameter.setRefType(RefType.INPUT);
        parameter.setRequired(true);

        startNode.addInputParameter(parameter);

        EndNode endNode = new EndNode();
        endNode.setId("endNode");

        Chain chain = new Chain();
        chain.addNode(startNode);
        chain.addNode(llmNode);
        chain.addNode(endNode);

        ChainEdge edge = new ChainEdge();
        edge.setId(UUID.randomUUID().toString());
        edge.setSource(startNode.getId());
        edge.setTarget(llmNode.getId());


        ChainEdge edge1 = new ChainEdge();
        edge1.setId(UUID.randomUUID().toString());
        edge1.setSource(llmNode.getId());
        edge1.setTarget(endNode.getId());

        chain.addEdge(edge);
        chain.addEdge(edge1);


        chain.addErrorListener(new ChainErrorListener() {
            @Override
            public void onError(Throwable throwable, Chain chain) {
                throwable.printStackTrace();
            }
        });
//
//
//        chain.addSuspendListener(new ChainSuspendListener() {
//
//            @Override
//            public void onSuspend(Chain chain) {
//                System.out.println("------------onSuspend-----------");
//                List<Parameter> suspendForParameters = chain.getSuspendForParameters();
//                System.out.println("suspendForParameters: " + suspendForParameters);
//
//                chain.resume(Maps.of("test","123"));
//            }
//        });


        chain.execute(Maps.of("input","你是谁"));
        System.out.printf("#########"+ chain.getMemory().get(edge.getId()));
    }
}
